"""
Authentication module for FedComm: pluggable backend, ECDSA and aggregate signature interface.
"""
import time
import hashlib
from abc import ABC, abstractmethod
from typing import Any, Tuple, List

# ECDSA backend
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature, decode_dss_signature
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.backends import default_backend

# Certless aggregate placeholder
from Crypto.PublicKey import ECC

class AuthScheme(ABC):
    @abstractmethod
    def keygen(self) -> Tuple[Any, Any]:
        pass

    @abstractmethod
    def sign(self, sk: Any, message: bytes, context: dict = None) -> Any:
        pass

    @abstractmethod
    def verify(self, pk: Any, message: bytes, signature: Any, context: dict = None) -> bool:
        pass

    @abstractmethod
    def aggregate_verify(self, list_of_pks: List[Any], messages: List[bytes], signatures: List[Any]) -> bool:
        pass

# --- ECDSA Implementation ---
class ECDSAAuth(AuthScheme):
    def keygen(self):
        sk = ec.generate_private_key(ec.SECP256R1(), default_backend())
        pk = sk.public_key()
        return sk, pk

    def sign(self, sk, message: bytes, context: dict = None):
        t0 = time.perf_counter()
        sig = sk.sign(message, ec.ECDSA(hashes.SHA256()))
        t1 = time.perf_counter()
        return sig, t1-t0

    def verify(self, pk, message: bytes, signature, context: dict = None):
        t0 = time.perf_counter()
        try:
            pk.verify(signature, message, ec.ECDSA(hashes.SHA256()))
            valid = True
        except Exception:
            valid = False
        t1 = time.perf_counter()
        return valid, t1-t0

    def aggregate_verify(self, list_of_pks, messages, signatures):
        # ECDSA 不支持聚合签名，逐个验证
        for pk, msg, sig in zip(list_of_pks, messages, signatures):
            if not self.verify(pk, msg, sig)[0]:
                return False
        return True

    @staticmethod
    def serialize_pk(pk) -> bytes:
        return pk.public_bytes(
            encoding=serialization.Encoding.DER,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )

    @staticmethod
    def deserialize_pk(data: bytes):
        return serialization.load_der_public_key(data, backend=default_backend())

    @staticmethod
    def serialize_signature(sig) -> bytes:
        return sig

    @staticmethod
    def deserialize_signature(data: bytes):
        return data

    @staticmethod
    def pk_size(pk) -> int:
        return len(ECDSAAuth.serialize_pk(pk))

    @staticmethod
    def signature_size(sig) -> int:
        return len(sig)

# --- Certless/Aggregate Placeholder ---
class CertlessAggregateAuth(AuthScheme):
    def keygen(self):
        # ECC keypair (placeholder, not pairing-based)
        sk = ECC.generate(curve='P-256')
        pk = sk.public_key()
        return sk, pk

    def sign(self, sk, message: bytes, context: dict = None):
        raise NotImplementedError("Aggregate/certless signature requires pairing-based crypto. Use charm-crypto for BLS/BGLS.")

    def verify(self, pk, message: bytes, signature, context: dict = None):
        raise NotImplementedError("Aggregate/certless signature requires pairing-based crypto. Use charm-crypto for BLS/BGLS.")

    def aggregate_verify(self, list_of_pks, messages, signatures):
        raise NotImplementedError("Aggregate/certless signature requires pairing-based crypto. Use charm-crypto for BLS/BGLS.")

    @staticmethod
    def pk_size(pk) -> int:
        return len(pk.export_key(format='DER'))

    @staticmethod
    def signature_size(sig) -> int:
        return len(sig)

# --- Helper functions ---
def hash_message(message: bytes) -> bytes:
    return hashlib.sha256(message).digest()

def serialize_signature(sig) -> bytes:
    return sig

def deserialize_signature(data: bytes):
    return data

def pk_size(pk) -> int:
    if hasattr(pk, 'public_bytes'):
        return ECDSAAuth.pk_size(pk)
    elif hasattr(pk, 'export_key'):
        return CertlessAggregateAuth.pk_size(pk)
    else:
        raise TypeError('Unknown pk type')

def signature_size(sig) -> int:
    return len(serialize_signature(sig))
